
import pickle
import geopandas as gpd
import pandas as pd
from shapely.geometry import Point
import warnings

from gpu_evaluator import GPUEvaluator


warnings.simplefilter(action='ignore', category=FutureWarning)


REAL_STORES_CSV = "福州市仓山区瑞幸咖啡坐标.csv"
ROAD_CACHE_PATH = "road_cache.pkl"
PSO_INPUT_DATA_PATH = "pso_input_data.pkl"


OUTPUT_SCORES_SHP = "luckin_real_store_scores.shp"
OUTPUT_SCORES_CSV = "luckin_real_store_scores.csv"

def load_csv_robustly(filepath: str) -> pd.DataFrame:
    """尝试用UTF-8和GBK两种编码加载CSV文件。"""
    try:
        # 优先尝试最标准的UTF-8
        return pd.read_csv(filepath, encoding='utf-8')
    except UnicodeDecodeError:
        print(f"  - 提示: 文件 '{filepath}' 使用UTF-8读取失败，正在尝试GBK...")
        # 如果UTF-8失败，再尝试GBK
        return pd.read_csv(filepath, encoding='gbk')
    except FileNotFoundError:
        # 如果文件不存在，直接抛出异常
        raise


def main():
    print("--- 开始对真实门店位置进行量化评估 ---")

    print("\n[步骤 1/4] 正在加载模型组件和基础数据...")
    try:
        with open(ROAD_CACHE_PATH, "rb") as f:
            road_network_nx = pickle.load(f)['full_graph']
        with open(PSO_INPUT_DATA_PATH, "rb") as f:
            pso_input_data = pickle.load(f)
        poi_data_3857 = pso_input_data["poi_data_3857"]
        population_points_3857 = pso_input_data["population_points_3857"]
        real_stores_df = load_csv_robustly(REAL_STORES_CSV)
        print("✅ 所有数据加载成功。")
    except FileNotFoundError as e:
        print(f"❌ 错误：找不到文件 '{e.filename}'。请确保所有必需文件都在场。")
        return
    except Exception as e:
        print(f"❌ 错误：加载文件失败。\n  详细错误: {e}")
        return

    print("\n[步骤 2/4] 正在初始化GPU评估器...")
    try:
        gpu_eval = GPUEvaluator(road_network_nx, population_points_3857, poi_data_3857)
    except Exception as e:
        print(f"❌ 错误：初始化GPUEvaluator失败。\n  详细错误: {e}")
        return

    print("\n[步骤 3/4] 正在处理并评估每一家真实门店...")
    try:
        # 🔥 [核心修改] 直接使用CSV中的坐标，不再调用 gcj2wgs
        locations_wgs84 = gpd.GeoDataFrame(
            real_stores_df,
            geometry=gpd.points_from_xy(real_stores_df.geolocation_lng, real_stores_df.geolocation_lat),
            crs="EPSG:4326"
        )
        locations_3857 = locations_wgs84.to_crs("EPSG:3857")
    except Exception as e:
        print(f"❌ 错误：处理门店坐标时失败。请检查CSV文件中的经纬度列名是否正确。\n  详细错误: {e}")
        return

    all_scores_data = []
    eval_config = {'pickup_radius': 300}
    for index, store in locations_3857.iterrows():
        store_coords = [(store.geometry.x, store.geometry.y)]
        fitness_scores = gpu_eval.calculate_luckin_fitness_final(store_coords, eval_config)
        single_store_score = fitness_scores.get('pickup_score', 0)
        store_name = store.get('name', f"门店_{index}")
        store_address = store.get('address', '')
        all_scores_data.append({
            'name': store_name,
            'address': store_address,
            'quality_score': single_store_score,
            'geometry': store.geometry
        })
        print(f"  - 已评估: {store_name} | 得分: {single_store_score:.2f}")

    print("\n[步骤 4/4] 正在将评估结果保存到文件...")
    if not all_scores_data:
        print("  - ⚠️ 警告：没有评估任何门店，不生成输出文件。")
        return
        
    scores_gdf = gpd.GeoDataFrame(all_scores_data, crs="EPSG:3857")
    try:
        scores_gdf.to_file(OUTPUT_SCORES_SHP, driver='ESRI Shapefile', encoding='utf-8')
        print(f"  - ✅ [可视化] 评估结果已保存到Shapefile: '{OUTPUT_SCORES_SHP}'")
    except Exception as e:
        print(f"  - ❌ 错误：保存Shapefile失败: {e}")

    scores_gdf_wgs84 = scores_gdf.to_crs("EPSG:4326")
    output_df = pd.DataFrame(scores_gdf_wgs84.drop(columns='geometry'))
    output_df['longitude'] = scores_gdf_wgs84.geometry.x
    output_df['latitude'] = scores_gdf_wgs84.geometry.y
    output_df = output_df.sort_values(by='quality_score', ascending=False).reset_index(drop=True)
    try:
        output_df.to_csv(OUTPUT_SCORES_CSV, index=False, encoding='utf-8-sig')
        print(f"  - ✅ [表格] 评估结果已保存到CSV: '{OUTPUT_SCORES_CSV}'")
    except Exception as e:
        print(f"  - ❌ 错误：保存CSV失败: {e}")

    print(f"\n🎉 [成功] 所有真实门店评估完成！")

if __name__ == "__main__":
    main()